#!/usr/bin/env python3.9
# coding=utf-8

"""
Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
Description: python api test.
Author: MindX SDK
Create: 2022
History: NA
"""

import logging
import time
from multiprocessing import JoinableQueue, Queue
import traceback
import numpy as np
import cv2
from mindx.sdk import base
from mindx.sdk.base import Image, ImageProcessor, Size, Tensor

from utils import car_post_process

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s',
                    filename="log.log", filemode='w')
logger = logging.getLogger(__name__)

# 车辆检测
class CarUnionDetection(object):
    def __init__(self, det_model_path, batch_size, device_id, pid):
        self.det_model = base.model(det_model_path, deviceId=device_id)
        self.image_processor = ImageProcessor(device_id) # initialize mxbase image_process
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid

        self.INFER_WAIT_TIME = 5

    def det_preprocess(self, image_ori, image_ori_height, image_ori_width):
        decodeImage = Image(image_ori[0], base.rgb)
        decodeImage.set_original_size(Size(image_ori_width, image_ori_height))
        image_ori = image_ori[:, :decodeImage.original_height, :decodeImage.original_width, :]
        image_ori = image_ori.transpose((0, 3, 1, 2))  # NCHW, RGB
        
        # image resized by dvpp
        decodeImage.to_device(self.device_id)
        resize_tuple, pad_tuple = car_post_process.letterbox(decodeImage)
        resize_conf = Size(resize_tuple[0], resize_tuple[1])
        decodeImage = self.image_processor.resize(decodeImage, resize_conf, base.bilinear_similar_opencv)
        decodeImage.to_host()

        # transfer to ndarray and put original and resized image array into queue
        image_src = np.array(decodeImage.to_tensor())  # NHWC
        image_src = image_src[:, :decodeImage.original_height, : decodeImage.original_width, :]
        image_src = cv2.copyMakeBorder(image_src[0], pad_tuple[0], pad_tuple[1], pad_tuple[2], pad_tuple[3],
                                       cv2.BORDER_CONSTANT, value=(112, 112, 112))
        image_src = np.expand_dims(image_src, axis=0).transpose((0, 3, 1, 2))  # NCHW
        return image_src, image_ori

    def infer(self, q_decode_car_union: Queue,
              q_rec_car_union_color: Queue,
              q_rec_car_union_prop: Queue,
              q_rec_car_union_brand: Queue,
              q_det_rec_carplate: Queue,
              q_track_car_union: Queue):
        try:
            logger.info("======================================== Start car_union detection inference ========================================")
            
            img_ndarray_list = []
            img_ori_ndarray_list = []

            count = 0
            while True:
                try:
                    _, image, image_ori_height, image_ori_width = q_decode_car_union.get(timeout=self.INFER_WAIT_TIME)
                except Exception as e:
                    if len(img_ndarray_list) > 0:
                        logger.info(f"{self.pid} pid process1_infer dynamic batch branch.")
                        current_num = len(img_ndarray_list)
                        pad_num = self.batch_size - current_num
                        pad_shape = (pad_num,) + img_ndarray_list[-1].shape
                        pad_zeros = np.zeros(shape=pad_shape)
                        img_ndarray_list = np.concatenate((np.array(img_ndarray_list), pad_zeros), axis=0)

                        # (1) Initialize one sample.
                        img_mxtensor = Tensor(np.squeeze(img_ndarray_list, axis=1))
                        img_mxtensor.to_device(self.device_id)
                        img_mxtensor_list = [img_mxtensor]

                        # (2) Det model infer
                        output_tensors = self.det_model.infer(img_mxtensor_list)  # output is a list with 4 arrays
                        for i, output_tensor in enumerate(output_tensors):
                            output_tensor.to_host()
                            output_tensors[i] = np.array(output_tensor)[:current_num]

                        # (3) post process
                        all_bboxes = car_post_process.det_postprocess(output_tensors[0], img_ndarray_list[:current_num], img_ori_ndarray_list)
                        rs = car_post_process.get_rs_from_box(all_bboxes, img_ori_ndarray_list)  # rs is the output

                        for i, res in enumerate(rs):
                            q_rec_car_union_color.put(res[0])
                            q_rec_car_union_prop.put(res[0])
                            q_rec_car_union_brand.put(res[0])
                            q_det_rec_carplate.put(res[0])
                            q_track_car_union.put(res)

                        # (4) Infer finished, free lists
                        img_ndarray_list = []
                        img_ori_ndarray_list = []
                    else:
                        logger.info(f"{self.pid} pid process1_infer wait time out, break.")
                        logger.info(f"{self.pid} pid process1_2_infer finished. get image cnt: {count}")
                        return
                else:
                    count += 1
                    img_src, img_ori = self.det_preprocess(image, image_ori_height, image_ori_width)

                    # 3 read input array and transfer array type, put into tensor list
                    img_ndarray = img_src.astype(np.float32)  # NCHW, RGB
                    img_ori_ndarray = img_ori  # NCHW, RGB
                    img_ndarray = img_ndarray / 255.
                    img_ndarray_list.append(img_ndarray)
                    img_ori_ndarray_list.append(img_ori_ndarray)

                    # 4. If BATCH_SIZE smaller than config, wait until get enough
                    if len(img_ndarray_list) < self.batch_size:
                        continue

                    # 5.1 Prepare batch input
                    img_mxtensor = Tensor(np.squeeze(np.array(img_ndarray_list), axis=1))
                    img_mxtensor.to_device(self.device_id)
                    img_mxtensor_list = [img_mxtensor]

                    # 5.2 Retina model infer
                    output_tensors = self.det_model.infer(img_mxtensor_list)  # output is a list with 4 arrays
                    for i, output_tensor in enumerate(output_tensors):
                        output_tensor.to_host()
                        output_tensors[i] = np.array(output_tensor)

                    # 5.3 Post process
                    all_bboxes = car_post_process.det_postprocess(output_tensors[0], img_ndarray_list, img_ori_ndarray_list)
                    rs = car_post_process.get_rs_from_box(all_bboxes, img_ori_ndarray_list)

                    for i, res in enumerate(rs):
                        q_rec_car_union_color.put(res[0])
                        q_rec_car_union_prop.put(res[0])
                        q_rec_car_union_brand.put(res[0])
                        q_det_rec_carplate.put(res[0])
                        q_track_car_union.put(res)

                    # 6 Clear lists and timer
                    img_ndarray_list = []
                    img_ori_ndarray_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid process1_infer failed")
            traceback.print_exc()

# 颜色识别
class CarUnionColorRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid

        self.INFER_WAIT_TIME = 10

    def rec_preprocess(self, img):
        height = 224
        width = 224
        # resize
        img = cv2.resize(img, (width, height))
        img = img[:, :, ::-1]  # BGR转RGB, 车牌颜色模型通道数在最后不用transpose
        img = img.astype(np.float32)

        # 归一化
        img /= 127.5
        img -= 1.
        return img

    def infer(self, q_get: Queue):
        try:
            logger.info("======================================== Start car_union rec color inference ========================================")
            count = 0
            img_arr_list = []

            while True:
                try:
                    img_input = q_get.get(timeout=self.INFER_WAIT_TIME)
                    if img_input is None:
                        continue
                except Exception as e:
                    if len(img_arr_list) > 0:
                        logger.info(f"{self.pid} pid car color process2_pattern_rec enter dynamic branch.")
                        current_num = len(img_arr_list)
                        pad_num = self.batch_size - current_num
                        pad_shape = (pad_num,) + img_arr_list[-1].shape
                        pad_zeros = np.zeros(shape=pad_shape)
                        img_arr_list = np.concatenate((np.array(img_arr_list), pad_zeros), axis=0)

                        # (1) Initialize one sample.
                        img_tensor = Tensor(img_arr_list)
                        img_tensor.to_device(self.device_id)
                        img_tensor_vec = [img_tensor]

                        # Rec model infer.
                        output_tensors = self.rec_model.infer(img_tensor_vec)
                        for i, output_tensor in enumerate(output_tensors):
                            output_tensor.to_host()
                            output_tensors[i] = np.array(output_tensor)[:current_num]
                        # Clear lists
                        img_arr_list = []
                    else:
                        logger.info(f"{self.pid} pid car color process2_pattern_rec wait time out, break.")
                        logger.info(f"{self.pid} pid car color process2_pattern_rec finished. get imagecnt: {count}")
                        return
                else:
                    img = img_input
                    count += 1

                    img = self.rec_preprocess(img)
                    img_arr_list.append(img)

                    if len(img_arr_list) < self.batch_size:
                        continue

                    tf_imgs = np.array(img_arr_list)
                    img_tensor = Tensor(tf_imgs)
                    img_tensor.to_device(self.device_id)
                    img_tensor_vec = [img_tensor]

                    # 4. rec model infer
                    output_tensors = self.rec_model.infer(img_tensor_vec)

                    # 5. rec model postprocess
                    for i, output_tensor in enumerate(output_tensors):
                        output_tensor.to_host()
                        output_tensors[i] = np.array(output_tensor)

                    ### process result
                    img_arr_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car color process2_pattern_rec failed")
            traceback.print_exc()

# 车辆类别
class CarUnionPropRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid, threshold=0.5):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid
        self.INFER_WAIT_TIME = 10
        self.threshold = threshold

    def rec_preprocess(self, img):
        height = 224
        width = 224

        # resize
        imh, imw = img.shape[:2]
        m = min(imh, imw)  # min dimension
        top, left = (imh - m) // 2, (imw - m) // 2
        img = cv2.resize(img[top:top + m, left:left + m], (width, height), interpolation=cv2.INTER_LINEAR)
        # img = cv2.resize(img, (width, height))

        img = img.transpose((2, 0, 1))  # BGR 转RGB
        img = img.astype(np.float32)

        # 归一化
        img /= 255.0

        # 标准化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        mean = np.array(mean).reshape(3, 1, 1)
        std = np.array(std).reshape(3, 1, 1)
        img = (img-mean)/std
        img = img.astype(np.float32)
        img = np.ascontiguousarray(img)

        return img

    def infer(self, q_get: Queue):
        def sigmoid(z):
            return 1/(1 + np.exp(-z))
        try:
            logger.info("======================================== Start car_union car prop inference ========================================")
            count = 0
            img_arr_list = []

            while True:
                try:
                    img_input = q_get.get(timeout=self.INFER_WAIT_TIME)
                    if img_input is None:
                        continue
                except Exception as e:
                    if len(img_arr_list) > 0:
                        current_num = len(img_arr_list)
                        pad_num = self.batch_size - current_num
                        pad_shape = (pad_num,) + img_arr_list[-1].shape
                        pad_zeros = np.zeros(shape=pad_shape)
                        img_arr_list = np.concatenate((np.array(img_arr_list), pad_zeros), axis=0)

                        # (1) Initialize one sample.
                        img_tensor = Tensor(img_arr_list)
                        img_tensor.to_device(self.device_id)
                        img_tensor_vec = [img_tensor]

                        # Rec model infer.
                        output_tensors = self.rec_model.infer(img_tensor_vec)
                        for i, output_tensor in enumerate(output_tensors):
                            output_tensor.to_host()
                            output_tensors[i] = np.array(output_tensor)[:current_num]
                        # Clear lists
                        img_arr_list = []
                    else:
                        logger.info(f"{self.pid} pid car prop process2_pattern_rec wait time out, break.")
                        logger.info(f"{self.pid} pid car prop process2_pattern_rec finished. get imagecnt: {count}")
                        return
                else:
                    img = img_input
                    count += 1
                    img = self.rec_preprocess(img)
                    img_arr_list.append(img)
                    if len(img_arr_list) < self.batch_size:
                        continue

                    tf_imgs = np.array(img_arr_list)
                    img_tensor = Tensor(tf_imgs)
                    img_tensor.to_device(self.device_id)
                    img_tensor_vec = [img_tensor]

                    # 4. rec model infer
                    output_tensors = self.rec_model.infer(img_tensor_vec)

                    # 5. rec model postprocess
                    for i, output_tensor in enumerate(output_tensors):
                        output_tensor.to_host()
                        output_tensors[i] = np.array(output_tensor)

                    ### process result
                    img_arr_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car prop process2_pattern_rec failed")
            traceback.print_exc()

# 车辆品牌
class CarUnionBrandRec(object):
    def __init__(self, rec_model_path, batch_size, device_id, pid, threshold=0.5):
        self.rec_model = base.model(rec_model_path, deviceId=device_id)
        self.batch_size = batch_size
        self.device_id = device_id
        self.pid = pid
        self.INFER_WAIT_TIME = 10
        self.threshold = threshold

    def rec_preprocess(self, img):
        height = 224
        width = 224

        # resize
        imh, imw = img.shape[:2]
        m = min(imh, imw)  # min dimension
        top, left = (imh - m) // 2, (imw - m) // 2
        img = cv2.resize(img[top:top + m, left:left + m], (width, height), interpolation=cv2.INTER_LINEAR)

        img = img.transpose((2, 0, 1))  # BGR
        img = img.astype(np.float32)

        # 归一化
        img /= 255.0

        # 标准化
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        mean = np.array(mean).reshape(3, 1, 1)
        std = np.array(std).reshape(3, 1, 1)
        img = (img-mean)/std
        img = img.astype(np.float32)
        img = np.ascontiguousarray(img)

        return img

    def infer(self, q_get: Queue):
        def sigmoid(z):
            return 1/(1 + np.exp(-z))
        try:
            logger.info("======================================== Start car_union car brand inference ========================================")
            count = 0
            img_arr_list = []

            while True:
                try:
                    img_input = q_get.get(timeout=self.INFER_WAIT_TIME)
                    if img_input is None:
                        continue
                except Exception as e:
                    if len(img_arr_list) > 0:
                        logger.info(f"{self.pid} pid car brand process2_pattern_rec enter dynamic branch.")
                        current_num = len(img_arr_list)
                        pad_num = self.batch_size - current_num
                        pad_shape = (pad_num,) + img_arr_list[-1].shape
                        pad_zeros = np.zeros(shape=pad_shape)
                        img_arr_list = np.concatenate((np.array(img_arr_list), pad_zeros), axis=0)

                        # (1) Initialize one sample.
                        img_tensor = Tensor(img_arr_list)
                        img_tensor.to_device(self.device_id)
                        img_tensor_vec = [img_tensor]

                        # Rec model infer.
                        output_tensors = self.rec_model.infer(img_tensor_vec)
                        for i, output_tensor in enumerate(output_tensors):
                            output_tensor.to_host()
                            output_tensors[i] = np.array(output_tensor)[:current_num]
                        # Clear lists
                        img_arr_list = []
                    else:
                        logger.info(f"{self.pid} pid car brand process2_pattern_rec wait time out, break.")
                        logger.info(f"{self.pid} pid car brand process2_pattern_rec finished. get image cnt: {count}")
                        return
                else:
                    img = img_input
                    count += 1
                    img = self.rec_preprocess(img)
                    img_arr_list.append(np.array(img))

                    if len(img_arr_list) < self.batch_size:
                        continue

                    img_tensor = Tensor(np.array(img_arr_list))
                    img_tensor.to_device(self.device_id)
                    img_tensor_vec = [img_tensor]

                    # 4. rec model infer
                    output_tensors = self.rec_model.infer(img_tensor_vec)

                    # 5. rec model postprocess
                    for i, output_tensor in enumerate(output_tensors):
                        output_tensor.to_host()
                        output_tensors[i] = np.array(output_tensor)

                    ### process result
                    img_arr_list = []
        except Exception as e:
            logger.error(f"{self.pid} pid car brand process2_pattern_rec failed")
            traceback.print_exc()


def car_union_detect(det_model_path, batch_size, device_id, pid,
                     q_decode_car_union: Queue,
                     q_rec_car_union_color: Queue,
                     q_rec_car_union_prop: Queue,
                     q_rec_car_union_brand: Queue,
                     q_det_rec_carplate: Queue,
                     q_track_car_union: Queue):
    CarUnionDetection(det_model_path, batch_size, device_id, pid).infer(
        q_decode_car_union,
        q_rec_car_union_color,
        q_rec_car_union_prop,
        q_rec_car_union_brand,
        q_det_rec_carplate,
        q_track_car_union)


def car_union_color_recognize(rec_model_path, batch_size, device_id, pid,
                              post_process_q_con: Queue):
    CarUnionColorRec(rec_model_path, batch_size, device_id, pid).infer(post_process_q_con)


def car_union_prop_recognize(rec_model_path, batch_size, device_id, pid,
                             post_process_q_con: Queue):
    CarUnionPropRec(rec_model_path, batch_size, device_id, pid).infer(post_process_q_con)


def car_union_brand_recognize(rec_model_path, batch_size, device_id, pid,
                              post_process_q_con: Queue):
    CarUnionBrandRec(rec_model_path, batch_size, device_id, pid).infer(post_process_q_con)
